
import numpy as np
import random
from itertools import groupby

class IndexGetter():
    def __init__(self,list_index_label,class_number=None,start=0):
        # list_index_label = list(enumerate(label_list))
        func_key = lambda i:i[1]
        list_index_label_sorted = sorted(list_index_label,key=func_key)
        group = groupby(list_index_label_sorted,key=func_key)
        self.index_list = [ list(map(lambda i:i[0],list(i))) for _,i in group]

        if class_number:
            assert(class_number==len(self.index_list))
        else:
            class_number = len(self.index_list)


        # # 数量非平均
        # len_list = list(map(len,self.index_list))
        # len_list.sort()
        # len_max = len_list[len(len_list)//3*2] # 按中位数长度切分
        # self.index_list = self.split(self.index_list,len_max)
        # random.shuffle(self.index_list)
        # class_number = len(self.index_list)

        self.index_now = [0]*class_number
        self.class_number = class_number
        self.start = start
        # print('len(label_list[0])',len(label_list[0]))
        # print('np.where(label_list==0).shape',np.where(label_list==0).shape)

    def get_index(self,label_target):
        """获得下一个 index_of_list"""
        # 取数
        index_of_list = self.index_now[label_target]
        index = self.index_list[label_target][index_of_list]

        # 超出范围 置0
        self.index_now[label_target]+=1
        if self.index_now[label_target]>=len(self.index_list[label_target]): # 其实只可能==
            self.index_now[label_target]=0
            np.random.shuffle(self.index_list[label_target])
        return index+self.start

    def get_index_random(self,label_target):
        list_target = self.index_list[label_target]
        index_of_list = random.randint(0,len(list_target)-1)
        index = list_target[index_of_list]
        return index+self.start

    def __len__(self):
        return self.class_number*max(map(lambda i:len(i),self.index_list))

    def split(self,index_list,len_max):
        result = []
        for list_ in index_list:
            count = max(1,len(list_)//len_max)
            len_temp = len(list_)//count
            for i in range(count-1):
                result.append(list_[i*len_temp:(i+1)*len_temp])
            result.append(list_[(count-1)*len_temp:])
        return result